import pandas as pd
import os
from openpyxl import load_workbook


class ExcelHandler:
    def __init__(self, file_path='../交易记录.xlsx'):
        self.file_path = file_path
        self._create_file_if_not_exists()

    def _create_file_if_not_exists(self):
        if not os.path.exists(self.file_path):
            df = pd.DataFrame(columns=['交易日期', '数量', '价格', '交易类型'])
            df.to_excel(self.file_path, sheet_name='示例品类', index=False)

    def get_categories(self):
        if os.path.exists(self.file_path):
            xls = pd.ExcelFile(self.file_path)
            return xls.sheet_names
        return []

    def save_record(self, category, date, quantity, price, trade_type):
        if not os.path.exists(self.file_path):
            self._create_file_if_not_exists()
        xls = pd.ExcelFile(self.file_path)
        try:
            df = xls.parse(category)
        except KeyError:
            df = pd.DataFrame(columns=['交易日期', '数量', '价格', '交易类型'])
        new_record = {
            '交易日期': date,
            '数量': quantity,
            '价格': price,
            '交易类型': trade_type
        }
        new_df = pd.DataFrame([new_record])
        df = pd.concat([df, new_df], ignore_index=True)
        with pd.ExcelWriter(self.file_path, mode='a', if_sheet_exists='replace') as writer:
            df.to_excel(writer, sheet_name=category, index=False)

    def get_latest_records(self, category):
        if os.path.exists(self.file_path):
            xls = pd.ExcelFile(self.file_path)
            try:
                df = xls.parse(category)
                return df.tail(10)
            except KeyError:
                return pd.DataFrame(columns=['交易日期', '数量', '价格', '交易类型'])
        return pd.DataFrame(columns=['交易日期', '数量', '价格', '交易类型'])

    def add_category(self, new_category):
        categories = self.get_categories()
        if new_category not in categories:
            df = pd.DataFrame(columns=['交易日期', '数量', '价格', '交易类型'])
            with pd.ExcelWriter(self.file_path, mode='a', if_sheet_exists='new') as writer:
                df.to_excel(writer, sheet_name=new_category, index=False)
            return True
        return False

    def rename_sheet(self, old_category, new_category):
        if os.path.exists(self.file_path):
            wb = load_workbook(self.file_path)
            if old_category in wb.sheetnames:
                if new_category not in wb.sheetnames:
                    ws = wb[old_category]
                    ws.title = new_category
                    wb.save(self.file_path)
                    return True
        return False